"""卫星和传感器模块 - 用于定义卫星和传感器对象"""

import numpy as np
from astropy import units as u
from astropy.time import Time, TimeDelta
from poliastro.bodies import Earth
from poliastro.twobody import Orbit
from poliastro.spacecraft import Spacecraft
from poliastro.earth import EarthSatellite
from enum import Enum
from typing import Tuple, Any

from propagator import OrbitPropagator, PropagationMethod
# 导入Observer抽象基类
from access import Observer, Observed

class SensorType(Enum):
    """传感器类型"""
    OPTICAL = "optical"           # 光学传感器
    RADAR = "radar"               # 雷达传感器
    INFRARED = "infrared"         # 红外传感器
    SAR = "sar"                   # 合成孔径雷达
    
class Sensor:
    """卫星传感器类"""
    
    def __init__(self, name, sensor_type, boresight, fov):
        """
        初始化传感器
        
        参数:
            name: 传感器名称
            sensor_type: 传感器类型
            boresight: 波束中心角(度), 距离卫星天顶点的角度
            fov: 视场角(度)
        """
        self.name = name
        self.sensor_type = sensor_type
        self.boresight = boresight * u.deg
        self.fov = fov * u.deg
        
    def __repr__(self):
        """字符串表示"""
        return (f"Sensor(name='{self.name}', type={self.sensor_type.value}, "
                f"boresight={self.boresight}, fov={self.fov})")
    
from astropy.coordinates import ITRS
class Satellite(Observer, Observed):
    """卫星类"""
    
    def __init__(self, name, orbit, mass=100*u.kg, area=1*u.m**2, C_D=2.2*u.one, sensors=None):
        """
        初始化卫星对象
        
        参数:
            name: 卫星名称
            orbit: 卫星轨道对象(poliastro Orbit)
            mass: 卫星质量
            area: 卫星横截面积
            C_D: 阻力系数(无量纲)
            sensors: 传感器列表
        """
        self.name = name
        self.orbit = orbit
        # 修复 Spacecraft 参数：需要 A(面积), C_D(阻力系数), m(质量)
        # 注意：Spacecraft 需要面积单位为 km²，而输入的 area 单位为 m²，需要转换
        area_km2 = area.to(u.km**2)
        self.spacecraft = Spacecraft(A=area_km2, C_D=C_D, m=mass)
        self.earth_satellite = EarthSatellite(orbit, self.spacecraft)
        self.sensors = sensors or []
        
        # 默认使用二体模型传播
        self.propagator = OrbitPropagator(orbit, PropagationMethod.KEPLER)
        # 预计算轨道数据
        self._precomputed_data = None
        
    def add_sensor(self, sensor):
        """
        添加传感器
        
        参数:
            sensor: 传感器对象
        """
        self.sensors.append(sensor)
        
    def set_propagation_method(self, method):
        """
        设置轨道递推方法
        
        参数:
            method: 递推方法
        """
        self.propagator = OrbitPropagator(self.orbit, method)
        
    def propagate(self, time_of_flight):
        """
        递推卫星轨道
        
        参数:
            time_of_flight: 递推时间
            
        返回:
            更新后的卫星对象
        """
        new_orbit = self.propagator.propagate(time_of_flight)
        
        # 创建新的卫星对象
        new_satellite = Satellite(
            self.name, 
            new_orbit, 
            self.spacecraft.m, 
            self.spacecraft.A,
            self.spacecraft.C_D,
            self.sensors
        )
        return new_satellite
    
    @staticmethod  
    def precompute_satellite_orbits(satellite, time_range: Tuple[Time, Time], step: TimeDelta):
        """
        预计算给定时间范围内该卫星的轨道数据
        
        参数:
            start_time: 开始时间
            end_time: 结束时间
            step: 时间步长
            
        返回:
            字典，包含时间和对应ITRS位置的数据
        """
        from access import AccessCalculator
        # 为卫星创建轨道数据字典
        satellite._precomputed_data = {
            'times': [],
            'itrs': []  # ITRS坐标系中的位置
        }
        
        start_time, end_time = time_range
        # 递推卫星轨道
        current_time = start_time
        while current_time <= end_time:
            prop_time = current_time - satellite.orbit.epoch
            current_orbit = satellite.propagator.propagate(prop_time)
            # 将卫星轨道转换为ITRS坐标
            satellite_itrs = AccessCalculator.transform_to_itrs(current_orbit)
            
            satellite._precomputed_data['times'].append(current_time)
            satellite._precomputed_data['itrs'].append(satellite_itrs)
            
            current_time += step
        return satellite, satellite._precomputed_data
        
    def get_precomputed_data(self, time_range:Tuple[Time, Time], step:TimeDelta):
        """获取预计算的轨道数据"""
        if self._precomputed_data is None:
            Satellite.precompute_satellite_orbits(self, time_range=time_range, step=step)
        return self._precomputed_data

    
    def find_itrs_from_precomputed_data(self, time:Time)->ITRS|None:
        """查找卫星的itrs坐标"""
        # 使用二分查找法查找最接近的时间
        import bisect
        index = bisect.bisect_left(self._precomputed_data['times'], time)
        if index < len(self._precomputed_data['times']):
            itrs = self._precomputed_data['itrs'][index]
        else:
            itrs = None
        return itrs

    def get_itrs_xyz(self, time:Time = None)->u.Quantity:
        """获取卫星的itrs坐标"""
        if time is None:
            time = Time.now()
        itrs_xyz = None
        if self._precomputed_data is not None:
            itrs = self.find_itrs_from_precomputed_data(time)
            itrs_xyz = itrs.cartesian.xyz
                
        if itrs_xyz is None:
            itrs_xyz = self.orbit.propagate(time - self.orbit.epoch).cartesian.xyz
        return itrs_xyz

    def can_see(self, target:Observed, time:Time = None) -> Tuple[bool]:
        """
        判断卫星是否可以看到目标
        
        参数:
            target_itrs: 观测目标(卫星或地面站)的itrs_xyz坐标
            
        返回:
            Tuple[bool]: (是否可见)
        """
        # 延迟导入，避免循环引用
        from access import AccessCalculator
        
        # 如果目标是卫星，暂时不支持卫星对卫星的可见性计算
        # TODO: 实现卫星对卫星的可见性计算
        if hasattr(target, 'orbit'):
            return False
            
        # 根据时间 self._precomputed_data['times'] 中找到最接近的时间
        if time is None:
            time = Time.now()
        itrs = self.find_itrs_from_precomputed_data(time)
        if itrs is not None:
            closest_itrs = itrs
        else:
            return False
        
        # 目标是地面站，使用卫星的传感器来确定可见性
        # 首先检查基本可见性 (仰角和地球遮挡)
        basic_visibility = target.can_see(self, time)
        #AccessCalculator.is_visible_on_the_ground(closest_itrs.cartesian.xyz, target.get_itrs_xyz(time))
        
        # 如果无法满足基本可见性条件，直接返回False
        if not basic_visibility:
            return False
            
        # 如果没有传感器，只考虑基本可见性
        if not self.sensors:
            return True
            
        # 如果有传感器，检查是否任一传感器可以观测到目标
        # TODO: 实现传感器对目标的可见性检查，考虑波束中心角(u.def)和视场角(u.deg)
        for sensor in self.sensors:
            # 计算目标相对于传感器波束中心的角度
            target_angle = AccessCalculator.calculate_angle(closest_itrs.cartesian.xyz, target.get_itrs_xyz(time))
            boresight = sensor.boresight
            fov = sensor.fov
            # 检查目标是否在传感器波束中心角和视场角之间
            if target_angle >= boresight - fov / 2 and target_angle <= boresight + fov / 2:
                return True

        return False
    
    def __getitem__(self, idx):
        """获取预计算的轨道数据"""
        return self._precomputed_data[idx]
    
    def __repr__(self):
        """字符串表示"""
        return (f"Satellite(name='{self.name}', "
                f"orbit={self.orbit}, "
                f"sensors={len(self.sensors)})")
                
class SatelliteConstellation:
    """卫星星座类"""
    
    def __init__(self, name, satellites=None):
        """
        初始化卫星星座
        
        参数:
            name: 星座名称
            satellites: 卫星列表
        """
        self.name = name
        self.satellites = satellites or []
        
    def add_satellite(self, satellite):
        """
        添加卫星到星座
        
        参数:
            satellite: 卫星对象
        """
        self.satellites.append(satellite)
        
    def propagate(self, time_of_flight):
        """
        递推星座中所有卫星
        
        参数:
            time_of_flight: 递推时间
            
        返回:
            更新后的星座对象
        """
        new_satellites = [sat.propagate(time_of_flight) for sat in self.satellites]
        
        # 创建新的星座对象
        new_constellation = SatelliteConstellation(self.name, new_satellites)
        
        return new_constellation
    
    def __len__(self):
        """返回星座中卫星数量"""
        return len(self.satellites)
        
    def __getitem__(self, idx):
        """获取索引对应的卫星"""
        return self.satellites[idx]
        
    def __repr__(self):
        """字符串表示"""
        return f"SatelliteConstellation(name='{self.name}', satellites={len(self.satellites)})" 